package cn.myeasyai.image;

import cn.myeasyai.config.MilvusConfig;
import cn.myeasyai.dto.ImageDTO;
import com.alibaba.fastjson.JSON;
import com.google.protobuf.Descriptors;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.FieldData;
import io.milvus.grpc.SearchResultData;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import io.swagger.models.auth.In;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;
import java.util.*;

@Service
public class ImageSearcher {
    @Resource
    private MilvusServiceClient milvusClient;
    @Resource
    private MilvusConfig milvusConfig;

    public List<ImageDTO> search(List<Float> floatList,Integer topK) {
        List<List<Float>> vectors = Collections.singletonList(floatList);
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(milvusConfig.getCollectionName()).withMetricType(MetricType.L2)
                .withTopK(topK)
                .withOutFields(Arrays.asList("image_url"))
                .withFloatVectors(vectors)
                .withVectorFieldName(milvusConfig.getVectorFieldName()).build();
        R<SearchResults> searchResults = this.milvusClient.search(searchParam);
        if (searchResults.getStatus() == 0) {
            final List<SearchResultsWrapper.IDScore> idScoreList = new SearchResultsWrapper(searchResults.getData().getResults()).getIDScore(0);
            List<ImageDTO> list = new ArrayList<>();
            idScoreList.forEach(idScore -> {
                ImageDTO imageDTO = new ImageDTO();
                final float score = idScore.getScore();
                final Map<String, Object> fieldValues = idScore.getFieldValues();
                imageDTO.setId(idScore.getLongID());
                imageDTO.setUrl(String.valueOf((fieldValues.getOrDefault("image_url", ""))));
                imageDTO.setScore(score);
                list.add(imageDTO);

            });
            return list;
        }
        return Collections.emptyList();
    }
}
